#!/bin/bash

# Copyright 2019 Mobvoi Inc. All Rights Reserved.
. ./path.sh || exit 1;

export CUDA_VISIBLE_DEVICES="0,1"
echo "CUDA_VISIBLE_DEVICES is ${CUDA_VISIBLE_DEVICES}"

stage=6 # start from 0 if you need to start from data preparation
stop_stage=6

# You should change the following two parameters for multiple machine training,
# see https://pytorch.org/docs/stable/elastic/run.html
HOST_NODE_ADDR="localhost:0"
num_nodes=1
job_id=2023


nj=16
dict=/home/work_nfs8/xlgeng/snsun/ckpt/fireredasr_xs/units.txt

data_type=shard_full_data

#train_config=/home/work_nfs8/xlgeng/snsun/code/wenet_firered/examples/aishell/firered/ckpt/fireredasr_xs/train.yaml
train_config=conf/train.yaml
dir=exp/firered_aed_xs_2
data=data/firered_aed_xs
mkdir -p $data
mkdir -p $dir
tensorboard_dir=tensorboard
#checkpoint=/home/work_nfs8/xlgeng/snsun/code/wenet_firered/examples/aishell/firered/exp/firered_aed_L/epoch_0.pt
#checkpoint=/home/work_nfs8/xlgeng/snsun/ckpt/fireredasr_xs/wenet_firered.pt
checkpoint=/home/work_nfs8/xlgeng/snsun/code/wenet_firered/examples/aishell/firered/exp/firered_aed_xs/step_269999.pt
num_workers=8
prefetch=10

train_data=$data/train.jsonl
cv_data=$data/cv.jsonl

data_conf=conf/data.yaml
python datahandle/handle_data_for_weight.py $data_conf $train_data
head -n 1 $train_data > $cv_data

# use average_checkpoint will get better result
average_checkpoint=true
decode_checkpoint=$dir/final.pt
average_num=3
decode_modes="attention"

train_engine=torch_ddp

deepspeed_config=conf/ds_stage2.json
deepspeed_save_states="model_only"

. tools/parse_options.sh || exit 1;


if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
  mkdir -p $dir
  num_gpus=$(echo $CUDA_VISIBLE_DEVICES | awk -F "," '{print NF}')
  dist_backend="nccl"
  if [ ${train_engine} == "deepspeed" ]; then
    echo "$0: using deepspeed"
  else
    echo "$0: using torch ddp"
  fi

  echo "$0: num_nodes is $num_nodes, proc_per_node is $num_gpus"
  torchrun --nnodes=$num_nodes --nproc_per_node=$num_gpus \
           --rdzv_id=$job_id --rdzv_backend="c10d" --rdzv_conf="join_timeout=3000" --rdzv_endpoint=$HOST_NODE_ADDR \
    wenet/bin/train.py \
      --train_engine ${train_engine} \
      --config $train_config \
      --data_type  $data_type \
      --train_data  $train_data \
      --cv_data $cv_data \
      ${checkpoint:+--checkpoint $checkpoint} \
      --model_dir $dir \
      --tensorboard_dir ${tensorboard_dir} \
      --ddp.dist_backend $dist_backend \
      --num_workers ${num_workers} \
      --prefetch ${prefetch} \
      --pin_memory \
      --deepspeed_config ${deepspeed_config} \
      --deepspeed.save_states ${deepspeed_save_states}
fi

if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
  # Test model, please specify the model you want to test by --checkpoint
  dir=/home/work_nfs8/xlgeng/snsun/code/wenet_firered/examples/aishell/firered/exp/firered_aed_xs
  average_checkpoint=false
  decode_checkpoint=/home/work_nfs8/xlgeng/snsun/code/wenet_firered/examples/aishell/firered/exp/firered_aed_xs/step_269999.pt
  if [ ${average_checkpoint} == true ]; then
    decode_checkpoint=$dir/avg_${average_num}.pt
    echo "do model average and final checkpoint is $decode_checkpoint"
    python wenet/bin/average_model.py \
      --dst_model $decode_checkpoint \
      --src_path $dir  \
      --num ${average_num} \
      --val_best
  fi
  # Please specify decoding_chunk_size for unified streaming and
  # non-streaming model. The default value is -1, which is full chunk
  # for non-streaming inference.
  decoding_chunk_size=
  ctc_weight=0.0
  reverse_weight=0.5
  python wenet/bin/recognize.py --gpu 0 \
    --modes $decode_modes \
    --config $dir/train.yaml \
    --data_type raw \
    --test_data /home/work_nfs8/xlgeng/snsun/child_data/data/test/scp/data.list \
    --checkpoint $decode_checkpoint \
    --beam_size 1 \
    --batch_size 1 \
    --blank_penalty 0.0 \
    --ctc_weight $ctc_weight \
    --reverse_weight $reverse_weight \
    --result_dir $dir \
    ${decoding_chunk_size:+--decoding_chunk_size $decoding_chunk_size}
  for mode in ${decode_modes}; do
    python tools/compute-wer.py --char=1 --v=1 \
      data/test/text $dir/$mode/text > $dir/$mode/wer
  done
fi

if [ ${stage} -le 6 ] && [ ${stop_stage} -ge 6 ]; then
  decode_modes="caption"
  dir=/home/work_nfs8/xlgeng/snsun/code/wenet_firered/examples/aishell/firered/exp/firered_aed_xs_2
  average_checkpoint=false
#  decode_checkpoint=/home/work_nfs8/xlgeng/snsun/code/wenet_firered/examples/aishell/firered/exp/firered_aed_xs/step_269999.pt
  decode_checkpoint=/home/work_nfs8/xlgeng/snsun/code/wenet_firered/examples/aishell/firered/exp/firered_aed_xs_2/epoch_1.pt
  decode_checkpoint=/home/work_nfs8/xlgeng/snsun/code/wenet_firered/examples/aishell/firered/exp/firered_aed_L/step_329999.pt
  decode_checkpoint=/home/work_nfs8/xlgeng/snsun/code/wenet_firered/examples/aishell/firered/exp/firered_aed_xs/epoch_0.pt
  if [ ${average_checkpoint} == true ]; then
    decode_checkpoint=$dir/avg_${average_num}.pt
    echo "do model average and final checkpoint is $decode_checkpoint"
    python wenet/bin/average_model.py \
      --dst_model $decode_checkpoint \
      --src_path $dir  \
      --num ${average_num} \
      --val_best
  fi
  # Please specify decoding_chunk_size for unified streaming and
  # non-streaming model. The default value is -1, which is full chunk
  # for non-streaming inference.
  decoding_chunk_size=
  ctc_weight=0.0
  reverse_weight=0.5
  python wenet/bin/recognize4caption.py --gpu 0 \
    --modes $decode_modes \
    --config $dir/train.yaml \
    --data_type raw \
    --test_data /home/work_nfs8/asr_data/data/test_sets_format_3000/caption_aslp_record/data.list \
    --checkpoint $decode_checkpoint \
    --beam_size 1 \
    --batch_size 12 \
    --blank_penalty 0.0 \
    --ctc_weight $ctc_weight \
    --reverse_weight $reverse_weight \
    --result_dir $dir \
    ${decoding_chunk_size:+--decoding_chunk_size $decoding_chunk_size}
  python handledata.py $dir/caption/text
#  for mode in ${decode_modes}; do
#    python tools/compute-wer.py --char=1 --v=1 \
#      data/test/text $dir/$mode/text > $dir/$mode/wer
#  done

fi

